import torch
import numpy as np

tartanair_class_names = [
    "empty",
    "ceiling",
    "floor",
    "wall",
    "window",
    "chair",
    "rug",
    "sofa",
    "screen",
    "tvs",
    "furn",
    "clock",
    "bonasai",
    "objs",
]
class_weights = torch.FloatTensor([0.05, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1])

class_freq_1_4 = np.array(
    [
        43744234,
        80205,
        1070052,
        905632,
        116952,
        180994,
        436852,
        279714,
        254611,
        28247,
        1805949,
        850724,
    ]
)
class_freq_1_8 = np.array(
    [
        5176253,
        17277,
        220105,
        183849,
        21827,
        33520,
        67022,
        44248,
        46615,
        4419,
        290218,
        142573,
    ]
)
class_freq_1_16 = np.array(
    [587620, 3820, 46836, 36256, 4241, 5978, 10939, 8000, 8224, 781, 49778, 25864]
)

tartanair_class_dict = {
    "empty": (0, [-1]),
    "ceiling": (1, [22, 147]),
    "floor": (2, [139]),
    "wall": (
        3,
        [90, 133, 144, 160, 172, 190, 193, 200, 208, 223, 224, 234, 244, 231, 239],
    ),
    "window": (4, [101, 146, 231]),
    "chair": (5, [211]),
    "rug": (6, [50, 207]),
    "sofa": (7, [120, 197]),
    "screen": (8, [125, 253]),
    "tvs": (
        9,
        [
            148,
            158,
        ],
    ),
    "furn": (10, [232, 173, 115, 144, 145, 152, 189, 173, 185, 205]),
    "clock": (11, [28]),
    "bonsai": (12, [137, 249]),
    "objs": (13, [-1]),
}
# BGR
class_colors = np.array(
    [
        [0, 0, 0],
        [153, 108, 6],
        [112, 105, 191],
        [89, 121, 72],
        [190, 225, 64],
        [206, 190, 59],
        [81, 13, 36],
        [115, 176, 195],
        [161, 171, 27],
        [135, 169, 180],
        [29, 26, 199],
        [102, 16, 239],
        [242, 107, 146],
        [156, 198, 23],
        [49, 89, 160],
        [68, 218, 116],
        [11, 236, 9],
        [196, 30, 8],
        [121, 67, 28],
        [0, 53, 65],
        [146, 52, 70],
        [226, 149, 143],
        [151, 126, 171],
        [194, 39, 7],
        [205, 120, 161],
        [212, 51, 60],
        [211, 80, 208],
        [189, 135, 188],
        [54, 72, 205],
        [103, 252, 157],
        [124, 21, 123],
        [19, 132, 69],
        [195, 237, 132],
        [94, 253, 175],
        [182, 251, 87],
        [90, 162, 242],
        [199, 29, 1],
        [254, 12, 229],
        [35, 196, 244],
        [220, 163, 49],
        [86, 254, 214],
        [152, 3, 129],
        [92, 31, 106],
        [207, 229, 90],
        [125, 75, 48],
        [98, 55, 74],
        [126, 129, 238],
        [222, 153, 109],
        [85, 152, 34],
        [173, 69, 31],
        [37, 128, 125],
        [58, 19, 33],
        [134, 57, 119],
        [218, 124, 115],
        [120, 0, 200],
        [225, 131, 92],
        [246, 90, 16],
        [51, 155, 241],
        [202, 97, 155],
        [184, 145, 182],
        [96, 232, 44],
        [133, 244, 133],
        [180, 191, 29],
        [1, 222, 192],
        [99, 242, 104],
        [91, 168, 219],
        [65, 54, 217],
        [148, 66, 130],
        [203, 102, 204],
        [216, 78, 75],
        [234, 20, 250],
        [109, 206, 24],
        [164, 194, 17],
        [157, 23, 236],
        [158, 114, 88],
        [245, 22, 110],
        [67, 17, 35],
        [181, 213, 93],
        [170, 179, 42],
        [52, 187, 148],
        [247, 200, 111],
        [25, 62, 174],
        [100, 25, 240],
        [191, 195, 144],
        [252, 36, 67],
        [241, 77, 149],
        [237, 33, 141],
        [119, 230, 85],
        [28, 34, 108],
        [78, 98, 254],
        [114, 161, 30],
        [75, 50, 243],
        [66, 226, 253],
        [46, 104, 76],
        [8, 234, 216],
        [15, 241, 102],
        [93, 14, 71],
        [192, 255, 193],
        [253, 41, 164],
        [24, 175, 120],
        [185, 243, 231],
        [169, 233, 97],
        [243, 215, 145],
        [72, 137, 21],
        [160, 113, 101],
        [214, 92, 13],
        [167, 140, 147],
        [101, 109, 181],
        [53, 118, 126],
        [3, 177, 32],
        [40, 63, 99],
        [186, 139, 153],
        [88, 207, 100],
        [71, 146, 227],
        [236, 38, 187],
        [215, 4, 215],
        [18, 211, 66],
        [113, 49, 134],
        [47, 42, 63],
        [219, 103, 127],
        [57, 240, 137],
        [227, 133, 211],
        [145, 71, 201],
        [217, 173, 183],
        [250, 40, 113],
        [208, 125, 68],
        [224, 186, 249],
        [69, 148, 46],
        [239, 85, 20],
        [108, 116, 224],
        [56, 214, 26],
        [179, 147, 43],
        [48, 188, 172],
        [221, 83, 47],
        [155, 166, 218],
        [62, 217, 189],
        [198, 180, 122],
        [201, 144, 169],
        [132, 2, 14],
        [128, 189, 114],
        [163, 227, 112],
        [45, 157, 177],
        [64, 86, 142],
        [118, 193, 163],
        [14, 32, 79],
        [200, 45, 170],
        [74, 81, 2],
        [59, 37, 212],
        [73, 35, 225],
        [95, 224, 39],
        [84, 170, 220],
        [159, 58, 173],
        [17, 91, 237],
        [31, 95, 84],
        [34, 201, 248],
        [63, 73, 209],
        [129, 235, 107],
        [231, 115, 40],
        [36, 74, 95],
        [238, 228, 154],
        [61, 212, 54],
        [13, 94, 165],
        [141, 174, 0],
        [140, 167, 255],
        [117, 93, 91],
        [183, 10, 186],
        [165, 28, 61],
        [144, 238, 194],
        [12, 158, 41],
        [76, 110, 234],
        [150, 9, 121],
        [142, 1, 246],
        [230, 136, 198],
        [5, 60, 233],
        [232, 250, 80],
        [143, 112, 56],
        [187, 70, 156],
        [2, 185, 62],
        [138, 223, 226],
        [122, 183, 222],
        [166, 245, 3],
        [175, 6, 140],
        [240, 59, 210],
        [248, 44, 10],
        [83, 82, 52],
        [223, 248, 167],
        [87, 15, 150],
        [111, 178, 117],
        [197, 84, 22],
        [235, 208, 124],
        [9, 76, 45],
        [176, 24, 50],
        [154, 159, 251],
        [149, 111, 207],
        [168, 231, 15],
        [209, 247, 202],
        [80, 205, 152],
        [178, 221, 213],
        [27, 8, 38],
        [244, 117, 51],
        [107, 68, 190],
        [23, 199, 139],
        [171, 88, 168],
        [136, 202, 58],
        [6, 46, 86],
        [105, 127, 176],
        [174, 249, 197],
        [172, 172, 138],
        [228, 142, 81],
        [7, 204, 185],
        [22, 61, 247],
        [233, 100, 78],
        [127, 65, 105],
        [33, 87, 158],
        [139, 156, 252],
        [42, 7, 136],
        [20, 99, 179],
        [79, 150, 223],
        [131, 182, 184],
        [110, 123, 37],
        [60, 138, 96],
        [210, 96, 94],
        [123, 48, 18],
        [137, 197, 162],
        [188, 18, 5],
        [39, 219, 151],
        [204, 143, 135],
        [249, 79, 73],
        [77, 64, 178],
        [41, 246, 77],
        [16, 154, 4],
        [116, 134, 19],
        [4, 122, 235],
        [177, 106, 230],
        [21, 119, 12],
        [104, 5, 98],
        [50, 130, 53],
        [30, 192, 25],
        [26, 165, 166],
        [10, 160, 82],
        [106, 43, 131],
        [44, 216, 103],
        [255, 101, 221],
        [32, 151, 196],
        [213, 220, 89],
        [70, 209, 228],
        [97, 184, 83],
        [82, 239, 232],
        [251, 164, 128],
        [193, 11, 245],
        [38, 27, 159],
        [229, 141, 203],
        [130, 56, 55],
        [147, 210, 11],
        [162, 203, 118],
        [255, 255, 255],
    ]
)
